package com.xs.langchain4j_springboot;

import com.xs.langchain4j_springboot.config.AiConfig;
import com.xs.langchain4j_springboot.splitter.CustomizedDocumentByLineSplitter;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.document.splitter.DocumentByRegexSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.context.annotation.Bean;

import java.net.URISyntaxException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

@SpringBootApplication(exclude = {DataSourceAutoConfiguration.class})
public class Langchain4jSpringbootApplication {

    public static void main(String[] args) {
        SpringApplication.run(Langchain4jSpringbootApplication.class, args);
    }

    /**
     * 对于一些必须要强制以一行的内容作为一个chunk的文档，可以用CustomizedDocumentByLineSplitter这个类，
     * 这个类extends了框架自带的DocumentByLineSplitter，删掉了一些原有的逻辑
     * 如果直接用框架自带的DocumentByLineSplitter的话，可能会导致多行的内容被放到一个chunk里的情况发生，比如category.txt这个文件
     * @param qwenEmbeddingModel
     * @param embeddingStore
     * @return
     * @throws URISyntaxException
     */
    @Bean
    CommandLineRunner getTermsOfServiceToVectorStore(QwenEmbeddingModel qwenEmbeddingModel, EmbeddingStore embeddingStore) throws URISyntaxException {
//        Path documentPath = Paths.get(Langchain4jSpringbootApplication.class.getClassLoader().getResource("rag/terms-of-service.txt").toURI());
        return args -> {
//            DocumentParser documentParser = new TextDocumentParser();
            Document document = ClassPathDocumentLoader.loadDocument("rag/category.txt");
//            DocumentByLineSplitter splitter = new DocumentByLineSplitter(15,5);
            CustomizedDocumentByLineSplitter splitter = new CustomizedDocumentByLineSplitter(20,1);
//            DocumentByLineSplitter splitter = new DocumentByLineSplitter(150,30);

            List<TextSegment> segmentList = splitter.split(document);
            List<Embedding> embeddings = qwenEmbeddingModel.embedAll(segmentList).content();
            embeddingStore.addAll(embeddings, segmentList);
        };
    }

}
